package com.egg.common.mybatis.interceptor;

import cn.hutool.core.annotation.AnnotationUtil;
import cn.hutool.core.util.ClassUtil;
import cn.hutool.core.util.ObjectUtil;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.extension.parser.JsqlParserSupport;
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
import com.egg.common.mybatis.annotation.DataScopeAnnotation;
import com.egg.common.mybatis.util.DataScopeUtils;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.LongValue;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.relational.ExistsExpression;
import net.sf.jsqlparser.statement.delete.Delete;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectBody;
import net.sf.jsqlparser.statement.select.SetOperationList;
import net.sf.jsqlparser.statement.update.Update;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;

import java.lang.reflect.Method;
import java.sql.Connection;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;


@SuppressWarnings({"unchecked"})
public class DataScopeInterceptor extends JsqlParserSupport implements InnerInterceptor {


    private final Map<String, DataScopeAnnotation> dataScopeAnnotationCacheMap = new ConcurrentHashMap<>();


    /**
     * 是否添加数据权限注解
     *
     * @param mappedStatementId
     * @return
     */
    private boolean hasDataScopeAnnotation(String mappedStatementId) {
        if (ObjectUtil.isEmpty(mappedStatementId)) {
            return false;
        }
        StringBuilder sb = new StringBuilder(mappedStatementId);
        int index = sb.lastIndexOf(".");
        String clazzName = sb.substring(0, index);
        Class<?> clazz = ClassUtil.loadClass(clazzName);
        String methodName = sb.substring(index + 1, sb.length());
        List<Method> methods = Arrays.stream(ClassUtil.getDeclaredMethods(clazz))
                .filter(method -> method.getName().equals(methodName)).collect(Collectors.toList());
        DataScopeAnnotation dataScopeAnnotation;
        for (Method method : methods) {
            dataScopeAnnotation = dataScopeAnnotationCacheMap.get(mappedStatementId);
            if (ObjectUtil.isNotNull(dataScopeAnnotation)) {
                return true;
            }
            if (AnnotationUtil.hasAnnotation(method, DataScopeAnnotation.class)) {
                dataScopeAnnotation = AnnotationUtil.getAnnotation(method, DataScopeAnnotation.class);
                dataScopeAnnotationCacheMap.put(mappedStatementId, dataScopeAnnotation);
                return true;
            }
        }
        return AnnotationUtil.hasAnnotation(clazz, DataScopeAnnotation.class);
    }

    /**
     * 需要追加的SQL类型
     *
     * @return
     */
    private List<SqlCommandType> sqlCommandTypeList() {
        return Arrays.asList(SqlCommandType.SELECT, SqlCommandType.UPDATE, SqlCommandType.DELETE);
    }

    /**
     * 是否放行
     *
     * @param ms
     * @return true-不追加数据权限, false-追加数据权限
     */
    private boolean passOnDataPerm(MappedStatement ms) {
        if (null == ms) {
            return false;
        }
        SqlCommandType sct = ms.getSqlCommandType();
        if (!sqlCommandTypeList().contains(sct)) {
            return true;
        }
        return !hasDataScopeAnnotation(ms.getId());
    }

    @Override
    public void beforePrepare(StatementHandler sh, Connection connection, Integer transactionTimeout) {
        PluginUtils.MPStatementHandler mpSh = PluginUtils.mpStatementHandler(sh);
        MappedStatement ms = mpSh.mappedStatement();
        if (passOnDataPerm(ms)) {
            return;
        }
        PluginUtils.MPBoundSql mpBs = mpSh.mPBoundSql();
        mpBs.sql(parserMulti(mpBs.sql(), ms.getId()));
    }

    @Override
    protected void processSelect(Select select, int index, String sql, Object obj) {
        if (select.getSelectBody() instanceof PlainSelect) {
            PlainSelect plainSelect = (PlainSelect) select.getSelectBody();
            Expression sqlSegment = getSqlSegment(plainSelect.getWhere(), plainSelect.getFromItem().toString());
            if (null != sqlSegment) {
                plainSelect.setWhere(sqlSegment);
            }
        } else if (select.getSelectBody() instanceof SetOperationList) {
            SetOperationList setOperationList = (SetOperationList) select.getSelectBody();
            List<SelectBody> selectBodyList = setOperationList.getSelects();
            selectBodyList.forEach(s -> {
                PlainSelect plainSelect = (PlainSelect) select.getSelectBody();
                Expression sqlSegment = getSqlSegment(plainSelect.getWhere(), plainSelect.getFromItem().toString());
                if (null != sqlSegment) {
                    plainSelect.setWhere(sqlSegment);
                }
            });
        }
    }

    @Override
    protected void processUpdate(Update update, int index, String sql, Object obj) {
        Expression sqlSegment = getSqlSegment(update.getWhere(), update.getFromItem().toString());
        if (null != sqlSegment) {
            update.setWhere(sqlSegment);
        }
    }

    @Override
    protected void processDelete(Delete delete, int index, String sql, Object obj) {
        Expression sqlSegment = getSqlSegment(delete.getWhere(), delete.getTable().getName());
        if (null != sqlSegment) {
            delete.setWhere(sqlSegment);
        }
    }

    /**
     * 处理有别名的主表名
     *
     * @param originalTableName
     * @return
     */
    private String parseTableName(String originalTableName) {
        // 使用空格分割字符串，获取第一个部分
        String[] parts = originalTableName.split("\\s+");
        return parts[0];
    }

    /**
     * 追加sql条件语句
     *
     * @param where     原sql条件
     * @param tableName 主表名
     * @return
     */
    private Expression getSqlSegment(Expression where, String tableName) {
        if (ObjectUtil.isEmpty(tableName)) {
            // 没有表名则直接返回,一般不可能出现该情况
            return null;
        }
        // 这里通过主表名获取相应的权限sql
        Optional<String> existSql = DataScopeUtils.buildDataScopeExistSql(parseTableName(tableName));
        if (!existSql.isPresent()) {
            return null;
        }
        ExistsExpression existsExpression = new ExistsExpression();
        existsExpression.setRightExpression(new LongValue(existSql.get()));
        if (null == where) {
            return existsExpression;
        } else {
            return new AndExpression(where, existsExpression);
        }
    }

}
